
import numpy as np
import json
import pandas
from tensorflow.keras.preprocessing import timeseries_dataset_from_array
import ast
from 数据处理 import 数组滑窗

def load(dataPath, batch_size, 方向, 解):
    print('---原始数据加载路径: ' + dataPath)

    pd_reader = pandas.read_csv(dataPath, sep=';', encoding='utf-8')

    if 方向 == "纵向":
        pd_reader = pd_reader.T

    # 跳过第一行时间 个股过多=窗口长度过大 会导致训练过程占用过多内存 可在此截取
    source = pd_reader.values[1:511]
    #  截取轴1(时间)防止数据量太大  可在此截取
    source = source[:, -15000:]
    total = source.shape[1]
    # 个股数量 = 窗口长度 , 因为模型需要正方形输入
    classy_len = source.shape[0]
    channel_len = len(json.loads(source[0][0]))

    # 遍历数组 字符串 使用json 转 array

     # 默认 float32 占用空间是 转化后的 float16 的两倍, 经试验有效减少内存占用, 但是会使标签失效
    # data = np.zeros([classy_len, total, channel_len], dtype=np.float16)
    data = np.zeros([classy_len, total, channel_len])
    for i in range(len(data)):
        for j in range(len(data[i])):
            data[i][j] = json.loads(source[i][j])

    # 使用keras的滑窗工具
    # 轴0应该是时间维度
    x_data = data.transpose(1,0,2)
    # 时间窗口(-window_len)前为输入  不需要这样 滑窗工具的x和y不一定要等长 相差一个窗口长度的数量
    # x_data = x_data[:-window_len]
    # 时间窗口(window_len)后列1(0:1)的2(2:3)通道为标签
    y_data = []
    if 解 == "线性":
        # 下一刻的最高价, 时间从classy_len开始, 第0:1个股, 第2:3通道
        y_data = x_data[classy_len:,0:1,2:3][:,0,0]
    # 简单起见, 直接比较最高点, 涨或平为1, 跌为0
    if 解 == "分类":
        # 下一刻的最高价, 时间从classy_len开始, 第0:1个股, 第2:3通道
        y_data_high = x_data[classy_len:, 1:2, 4:5][:, 0, 0]
        # 最后时刻的收盘价, 时间从classy_len-1开始, 第0:1个股, 第4:5通道
        y_data_close = x_data[classy_len-1:, 1:2, 4:5][:, 0, 0]
        for i in range(len(y_data_high)):
            if y_data_high[i] > y_data_close[i]:
                y_data.append(1)
            else:
                y_data.append(0)

    # start_index和end_index 好像又bug  手动切分训练和验证集
    # 验证集起始索引
    valid_index = -1000
    train_dataset = timeseries_dataset_from_array(data=x_data[:valid_index], targets=y_data[:valid_index], sequence_length=classy_len, shuffle=True, sequence_stride=1, sampling_rate=1, batch_size=batch_size)
    # 倒数的时候 标签y要后退窗口长度 才能和x对应上
    validate_dataset = timeseries_dataset_from_array(data=x_data[valid_index:], targets=y_data[valid_index+classy_len:], sequence_length=classy_len,  sequence_stride=1, sampling_rate=1, batch_size=batch_size)
    # 验证滑窗数据正确性时不能混洗
    for batch in train_dataset:
        inputs, targets = batch
        inputs.numpy()
        xx = np.array_equal(inputs[0], x_data[:classy_len])
        yy = np.array_equal(targets[0], y_data[0])
        break
    for batch in validate_dataset:
        inputs, targets = batch
        inputs.numpy()
        xx = np.array_equal(inputs[0], x_data[valid_index:classy_len+valid_index])
        yy = np.array_equal(targets[0], y_data[valid_index+classy_len])
        break
    shape = [classy_len, classy_len, channel_len]
    return train_dataset, validate_dataset, shape, np.max(y_data) + 1



if __name__ == '__main__':
    load("../temp/data/通信达/分钟线.csv", 32, "纵向", "分类")
    # load("纵向数组测试数据.csv", 32, "纵向", "分类")
    # load("横向数组测试数据.csv", 32, "横向', "分类")
